from chernc.treediff.gumtree.matchers.mapping_store import MappingStore
from chernc.treediff.gumtree.matchers.matcher import Matcher
from chernc.treediff.gumtree.tree.tree import TreeLike
from chernc.treediff.gumtree.tree.tree_utils import TreeUtils


class IdMatcher(Matcher):
    def match(self, src: TreeLike, dst: TreeLike, mappings: MappingStore) -> MappingStore:
        # Create candidate mappings for the source tree
        src_candidate_mappings = {}
        for t in TreeUtils.pre_order(src):
            id = t.get_metadata("id")
            if id is not None:
                if id not in src_candidate_mappings:
                    src_candidate_mappings[id] = set()
                src_candidate_mappings[id].add(t)

        # Create candidate mappings for the destination tree
        dst_candidate_mappings = {}
        for t in TreeUtils.pre_order(dst):
            id = t.get_metadata("id")
            if id is not None:
                if id not in dst_candidate_mappings:
                    dst_candidate_mappings[id] = set()
                dst_candidate_mappings[id].add(t)

        # Perform the matching based on IDs
        for id, src_set in src_candidate_mappings.items():
            if len(src_set) == 1 and id in dst_candidate_mappings and len(dst_candidate_mappings[id]) == 1:
                mappings.add_mapping(next(iter(src_set)), next(iter(dst_candidate_mappings[id])))

        return mappings
